package pl.edu.pb.wi.pwnography.modules;

import java.io.ByteArrayOutputStream;
import java.io.PrintStream;
import java.io.Serializable;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.logging.Level;
import java.util.logging.Logger;

import pl.edu.pb.wi.pwnography.model.ParsedData;
import quickdt.Attributes;
import quickdt.Instance;
import quickdt.Leaf;
import quickdt.Node;
import quickdt.TreeBuilder;

/*
 * Moduł oparty o bibliotekę QuickDT.
 * 
 */
public class DecisionTree {
    private static final Logger log = Logger.getLogger(DecisionTree.class
	    .getName());
    protected static final String PREFIX_BASE = "DT_";
    protected static final String CLASSES_SUFFIX = "_CLASS";

    public static List<Object> leaveOneOut(ParsedData data,
	    List<String> excludedColumns, String classColumn,
	    int discretizationLevel, int rowCount,
	    Map<String, List<Object>> rows) throws Exception {
	if (discretizationLevel < 2 || discretizationLevel > 5)
	    throw new Exception(
		    "Discretization level should be between 2 and 5.");
	if (!data.getColumnNames().contains(classColumn))
	    throw new Exception("Class column defined by user doesn't exist.");

	List<String> dataColumns = new ArrayList<String>(
		data.getFloatTypeColumnNames());
	if (excludedColumns != null) {
	    dataColumns.removeAll(excludedColumns);
	}
	if (dataColumns.contains(classColumn)) {
	    dataColumns.remove(classColumn);
	}

	List<String> allColumns = new ArrayList<String>(
		data.getStringTypeColumnNames());
	allColumns.addAll(dataColumns);
	if (excludedColumns != null) {
	    allColumns.removeAll(excludedColumns);
	}
	if (allColumns.contains(classColumn)) {
	    allColumns.remove(classColumn);
	}

	if (allColumns.size() == 0)
	    throw new Exception("You can't exclude all columns.");

	/*
	 * Dyskretyzacja atrybutów ciągłych.
	 */
	Map<String, List<Object>> namedDiscretizedColumn = new LinkedHashMap<String, List<Object>>();
	for (String colName : dataColumns) {
	    namedDiscretizedColumn.put(colName, Discretization.largestClass(
		    data.getColumnValues(colName), discretizationLevel,
		    new LinkedHashMap<Object, Integer>()));
	}

	List<Object> classesBasedOnTree = new ArrayList<Object>();
	for (int i = 0; i < rowCount; ++i) {
	    /*
	     * Budujemy drzewko.
	     */
	    boolean lock = true;
	    Set<Instance> instances = new HashSet<Instance>();
	    for (int j = 0; j < rowCount; ++j) {
		if (i == j)
		    continue;

		List<Serializable> attributes = new ArrayList<Serializable>();
		for (String colName : allColumns) {
		    attributes.add(colName);

		    Object val;
		    if (namedDiscretizedColumn.containsKey(colName))
			val = namedDiscretizedColumn.get(colName).get(j);
		    else
			val = rows.get(colName).get(j);

		    attributes.add((val instanceof Float) ? ((Float) val)
			    : ((String) val));
		}

		Object clazz = rows.get(classColumn).get(j);

		instances.add(Attributes.create(
			(Serializable[]) attributes
				.toArray(new Serializable[attributes.size()]))
			.classification(
				(clazz instanceof Float) ? ((Float) clazz)
					: ((String) clazz)));
	    }

	    TreeBuilder treeBuilder = new TreeBuilder();
	    Node tree = treeBuilder.buildTree(instances);

	    List<Serializable> attributes = new ArrayList<Serializable>();
	    for (String colName : allColumns) {
		attributes.add(colName);

		Object val;
		if (namedDiscretizedColumn.containsKey(colName))
		    val = namedDiscretizedColumn.get(colName).get(i);
		else
		    val = rows.get(colName).get(i);

		attributes.add((val instanceof Float) ? ((Float) val)
			: ((String) val));
	    }

	    Leaf leaf = tree.getLeaf(Attributes
		    .create((Serializable[]) attributes
			    .toArray(new Serializable[attributes.size()])));
	    classesBasedOnTree.add(leaf.classification.toString());
	    if (lock && log.isLoggable(Level.INFO)) {
		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		tree.dump(new PrintStream(baos));
		log.info(new String(baos.toByteArray()));
		log.info("\n============================================\n");

		lock = false;
	    }
	}

	return classesBasedOnTree;
    }

    public static Map<String, List<Object>> createColumn(String classColumn,
	    List<Object> values, int discretizationLevel) {

	/*
	 * Tworzenie nazwy nowej kolumny oraz umieszczenie jej razem z
	 * wartościami w HashMapie.
	 */
	String prefixedName = new StringBuilder(PREFIX_BASE)
		.append(String.format("%d_DISCRETIZATION_LEVEL",
			discretizationLevel)).append("_").append(classColumn)
		.append(CLASSES_SUFFIX).toString();

	Map<String, List<Object>> namedDtColumn = new HashMap<String, List<Object>>();
	namedDtColumn.put(prefixedName, values);

	return namedDtColumn;
    }
}
